import argparse
import json
import matplotlib.pyplot as plt
import sys
from pathlib import Path


def plot_experiment_results(results_file: str, output_file: str):
    """
    绘制实验结果折线图
    
    Args:
        results_file: 实验结果 JSON 文件路径
        output_file: 输出图片文件路径
    """
    # 读取实验结果
    try:
        with open(results_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except FileNotFoundError:
        print(f"错误: 找不到结果文件 {results_file}")
        sys.exit(1)
    except json.JSONDecodeError:
        print(f"错误: 无法解析 JSON 文件 {results_file}")
        sys.exit(1)
    
    # 提取数据
    accuracy_results = data.get('accuracy_results', {})
    
    if not accuracy_results:
        print("错误: 结果文件中没有准确率数据")
        sys.exit(1)
    
    # 准备绘图数据
    methods = ['idf', 'bm25', 'entropy']
    topks = [1, 3, 5]
    colors = ['#2E86AB', '#F18F01', '#A23B72']
    markers = ['o', 's', '^']
    
    # 设置中文字体支持
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 创建图形
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # 绘制每种方法的折线
    for i, method in enumerate(methods):
        if method not in accuracy_results:
            print(f"警告: 结果中没有方法 {method} 的数据")
            continue
        
        accuracies = []
        for topk in topks:
            acc = accuracy_results[method].get(str(topk), 0.0)
            accuracies.append(acc * 100)  # 转换为百分比
        
        # 绘制折线
        ax.plot(topks, accuracies, 
                color=colors[i], 
                marker=markers[i], 
                linewidth=2.5,
                markersize=8,
                label=method.upper(),
                alpha=0.8)
        
        # 在数据点上添加数值标签
        for j, (topk, acc) in enumerate(zip(topks, accuracies)):
            ax.text(topk, acc + 0.5, f'{acc:.1f}%', 
                   ha='center', va='bottom', fontsize=9)
    
    # 设置图表属性
    ax.set_xlabel('Top-K 参数', fontsize=12, fontweight='bold')
    ax.set_ylabel('准确率 (%)', fontsize=12, fontweight='bold')
    ax.set_title('不同权重方法在不同 Top-K 参数下的准确率对比', 
                fontsize=14, fontweight='bold', pad=20)
    ax.set_xticks(topks)
    ax.set_xticklabels([f'Top-{k}' for k in topks])
    ax.legend(loc='best', fontsize=11, framealpha=0.9)
    ax.grid(True, alpha=0.3, linestyle='--')
    ax.set_ylim(0, 100)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图片
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"折线图已保存为 {output_file}")
    
    # 显示图表
    plt.show()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='绘制实验结果折线图')
    parser.add_argument('--results_file', type=str,
                        default='data/experiment_results.json',
                        help='实验结果 JSON 文件路径')
    parser.add_argument('--output_file', type=str,
                        default='./data/weight_method_comparison.png',
                        help='输出图片文件路径')
    
    args = parser.parse_args()
    
    # 绘制图表
    plot_experiment_results(args.results_file, args.output_file)


if __name__ == '__main__':
    main()
